package com.hub.realtime.flink.submit.base;


import com.hub.realtime.common.exception.UtilException;
import com.hub.realtime.common.model.FlinkRunRequest;
import com.hub.realtime.common.model.FlinkStopRequest;
import com.hub.realtime.common.model.FlinkSubmitResponse;
import com.hub.realtime.common.utils.ListUtils;
import com.hub.realtime.flink.submit.model.FlinkSubmitCommon;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.Options;
import org.apache.commons.lang3.StringUtils;
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.JobID;
import org.apache.flink.client.cli.CliArgsException;
import org.apache.flink.client.cli.CliFrontend;
import org.apache.flink.client.cli.CliFrontendParser;
import org.apache.flink.client.cli.CustomCommandLine;
import org.apache.flink.configuration.*;
import org.apache.flink.runtime.jobgraph.SavepointConfigOptions;
import org.apache.flink.runtime.security.SecurityConfiguration;
import org.apache.flink.runtime.security.SecurityUtils;
import org.apache.flink.util.FlinkException;
import org.apache.flink.yarn.configuration.YarnDeploymentTarget;

import java.net.MalformedURLException;
import java.net.URL;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import static org.apache.flink.configuration.CoreOptions.*;


@Slf4j
public abstract class FlinkSubmitBase {
    public abstract FlinkSubmitResponse submit(FlinkRunRequest flinkRunRequest);

    public abstract String stop(FlinkStopRequest flinkStopRequest) throws FlinkException;

    /**
     * 获取flink的默认配置
     *
     * @param flinkConfDir
     * @return
     */
    protected Configuration getDefaultFLinkConf(String flinkConfDir) {
        return GlobalConfiguration.loadConfiguration(flinkConfDir);
    }

    /**
     * 转换jobId
     *
     * @param jobIdString
     * @return
     */
    public JobID parseJobId(String jobIdString) {
        if (StringUtils.isEmpty(jobIdString)) throw new UtilException("任务的job id不能为空");
        final JobID jobId;
        try {
            jobId = JobID.fromHexString(jobIdString);
        } catch (IllegalArgumentException e) {
            throw new UtilException(e.getMessage());
        }
        return jobId;
    }


    /**
     * 初始化公共部分
     *
     * @param flinkRunRequest
     * @return
     */
    protected FlinkSubmitCommon initalCommon(FlinkRunRequest flinkRunRequest) {

        FlinkSubmitCommon flinkSubmitCommon = new FlinkSubmitCommon();
        String configDir = flinkRunRequest.getFlinkHome().concat("/conf");
        Configuration defaultConfig = getDefaultFLinkConf(configDir);
        flinkSubmitCommon.setDefaultConfig(defaultConfig);
        List<CustomCommandLine> customCommandLines = CliFrontend.loadCustomCommandLines(defaultConfig, configDir);
        CliFrontend cli = new CliFrontend(defaultConfig, customCommandLines);
        try {
            SecurityUtils.install(new SecurityConfiguration(cli.getConfiguration()));
            flinkSubmitCommon.setCliFrontend(cli);
            Options commandOptions = CliFrontendParser.getRunCommandOptions();
            flinkSubmitCommon.setBaseCommandOptions(commandOptions);
            CommandLine effectiveCommandLine = getEffectiveCommandLine(flinkRunRequest, flinkSubmitCommon);
            flinkSubmitCommon.setCommandLine(effectiveCommandLine);
            CustomCommandLine activeCommandLine = cli.validateAndGetActiveCommandLine(effectiveCommandLine);
            flinkSubmitCommon.setActiveCommandLine(activeCommandLine);
            return flinkSubmitCommon;
        } catch (Exception e) {
            log.error("初始化公共部分出错：" + e.getMessage());
            e.printStackTrace();
            throw new UtilException("初始化公共部分出错：" + e.getMessage());
        }
    }

    /**
     * 通用的flink配置
     *
     * @param flinkRunRequest
     * @return
     */
    protected <T> Configuration applyToConfiguration(FlinkRunRequest flinkRunRequest, List<T> jobJars) {
        Configuration configuration = new Configuration();
        //设置并行度
        if (flinkRunRequest.getParallelism() != ExecutionConfig.PARALLELISM_DEFAULT) {
            configuration.setInteger(CoreOptions.DEFAULT_PARALLELISM, flinkRunRequest.getParallelism());
        }
        //设置分离模式
        configuration.setBoolean(DeploymentOptions.ATTACHED, false);
        if (flinkRunRequest.getExecuteMode().equals(YarnDeploymentTarget.PER_JOB.getName())) {
            configuration.setBoolean(
                    DeploymentOptions.SHUTDOWN_IF_ATTACHED, true);
        }
        if (flinkRunRequest.getSavePointConf() != null) {
            if (flinkRunRequest.getSavePointConf().getAllowNonRestoredState()) {
                configuration.setBoolean(
                        SavepointConfigOptions.SAVEPOINT_IGNORE_UNCLAIMED_STATE,
                        flinkRunRequest.getSavePointConf().getAllowNonRestoredState());
            }
            if (StringUtils.isNotEmpty(flinkRunRequest.getSavePointConf().getSavePointDir())) {
                configuration.setString(SavepointConfigOptions.SAVEPOINT_PATH, flinkRunRequest.getSavePointConf().getSavePointDir());
            }
        }
        //动态参数
        if (flinkRunRequest.getOptions() != null && flinkRunRequest.getOptions().size() > 0) {
            for (Map.Entry<String, String> op : flinkRunRequest.getOptions().entrySet()) {
                if (StringUtils.isNotEmpty(op.getValue())) {
                    String key = op.getKey();
                    if (key.equals("env.java.opts")) {
                        configuration.setString(FLINK_JVM_OPTIONS, op.getValue());
                    } else if ("env.java.opts.jobmanager".equals(op.getKey())) {
                        configuration.setString(FLINK_JM_JVM_OPTIONS, op.getValue());
                    } else if ("env.java.opts.taskmanager".equals(op.getKey())) {
                        configuration.setString(FLINK_TM_JVM_OPTIONS, op.getValue());
                    } else if ("env.java.opts.historyserver".equals(op.getKey())) {
                        configuration.setString(FLINK_HS_JVM_OPTIONS, op.getValue());
                    } else if ("env.java.opts.client".equals(op.getKey())) {
                        configuration.setString(FLINK_CLI_JVM_OPTIONS, op.getValue());
                    } else {
                        configuration.setString(op.getKey(), op.getValue());
                    }
                }
            }
        }
        ConfigUtils.encodeCollectionToConfig(
                configuration, PipelineOptions.JARS, jobJars, Object::toString);

        //设置类加载顺序
        if (StringUtils.isNotEmpty(flinkRunRequest.getResolveOrder())) {
            configuration.set(CoreOptions.CLASSLOADER_RESOLVE_ORDER, flinkRunRequest.getResolveOrder());
        }

        return configuration;
    }


    /**
     * 获取当前的命令行工具
     *
     * @param runRequest
     * @return
     */
    @SneakyThrows
    private CommandLine getEffectiveCommandLine(FlinkRunRequest runRequest,
                                                FlinkSubmitCommon flinkSubmitCommon) {
        List<String> runArgs = new ArrayList();
        if (StringUtils.isEmpty(runRequest.getExecuteMode()))
            throw new UtilException("请选择任务运行模式：yarn-application|yarn-per-job");
        if (runRequest.getExecuteMode().contains("yarn")) {
            if (runRequest.getExecuteMode().equals("yarn-application")) {
                runArgs.add("-t" + YarnDeploymentTarget.APPLICATION.getName());

            } else {
                runArgs.add("-t" + YarnDeploymentTarget.PER_JOB.getName());
                runArgs.add("-m" + "yarn-cluster");
                runArgs.add("-c" + runRequest.getMainClass());
            }
        }
        Options commandLineOptions =
                CliFrontendParser.mergeOptions(flinkSubmitCommon.getBaseCommandOptions(),
                        flinkSubmitCommon.getCliFrontend().getCustomCommandLineOptions());
        String[] args = ListUtils.toArray(runArgs, String.class);
        try {
            return CliFrontendParser.parse(commandLineOptions, args, true);
        } catch (CliArgsException e) {
            log.error("获取命令行参数工具失败：" + e.getMessage());
            e.printStackTrace();
            throw new UtilException("获取命令行参数工具失败：" + e.getMessage());
        }
    }

    /**
     * 获取第三方依赖
     *
     * @param flinkRunRequest
     * @return
     */
    @SneakyThrows
    public List<URL> getClassPaths(FlinkRunRequest flinkRunRequest) {
        List<URL> classPaths = new ArrayList<>();
        flinkRunRequest.getDependencies().forEach(d -> {
            try {
                classPaths.add(new URL(d));
            } catch (MalformedURLException e) {
                e.printStackTrace();
                throw new UtilException("Bad syntax for classpath: " + d);
            }
        });
        return classPaths;
    }


}
